import os
import json
import torch
import warnings
import argparse
import numpy as np
from collections import defaultdict
from typing import Dict
from torch.utils.data import DataLoader, Dataset
from sklearn.metrics import accuracy_score, precision_score, recall_score, roc_auc_score

from train_eval_fund_2 import TrainerDR
from train_eval_fund_4 import DeepSurModel
from config import TrainerConfig

device = 'cuda:1' if torch.cuda.is_available() else 'cpu'
result_cache = defaultdict(list)
class_name = ['正常_N', '糖尿病_D', '青光眼_G', '白内障_C', 'AMD_A', '高血压_H', '近视_M', '其他_O']


def batch(model, data) -> dict:
    imgs = data['image'].to(device)
    labels = data['labels'].to(device).float()
    
    labels = labels.squeeze(1)
    # predictions,_ = model(imgs)
    binary_logits, disease_logits= model(imgs)
    predictions = hierarchical_predict(binary_logits, disease_logits)
    return dict(
        predictions=predictions,
        labels=labels,
    )

# 正式测试用
def batch_formal(model, data) -> dict:
    left_imgs = data['left_image'].to(device)
    right_imgs = data['right_image'].to(device)
    labels = data['labels'].to(device)

    labels = labels.squeeze(1)
    # left_pred,_ = model(left_imgs)
    # right_pred,_ = model(right_imgs)
    binary_logits_l, disease_logits_l = model(left_imgs)
    left_pred = hierarchical_predict(binary_logits_l, disease_logits_l)
    binary_logits_r, disease_logits_r = model(left_imgs)
    right_pred = hierarchical_predict(binary_logits_r, disease_logits_r)
    predictions = merge_predicitons(left_pred, right_pred)
    return dict(
        predictions=predictions,
        labels=labels
    )

def merge_predicitons(left, right):
    left = left.to(torch.int64)
    right = right.to(torch.int64)
    N = left[:, 0] & right[:, 0]
    rest = left[:, 1:] | right[:, 1:]
    
    pred = torch.cat((N.unsqueeze(1), rest), dim=1)
    return pred

def hierarchical_predict(binary_pred, disease_pred, threshold=0.5):
    binary_pred = torch.tensor(binary_pred, dtype=torch.float32)
    disease_pred = torch.tensor(disease_pred, dtype=torch.float32)

    final_pred = torch.zeros(binary_pred.shape[0], 8, device=device)  # 8分类结果
    # 正常样本标记为类别0
    normal_mask = (binary_pred >= threshold).squeeze()
    final_pred[normal_mask, 0] = 1

    # 异常样本选择疾病的类别
    abnormal_mask = ~normal_mask
    if abnormal_mask.sum() > 0:
        disease_labels = (disease_pred[abnormal_mask] >= threshold).float()  # 转换为0/1
        final_pred[abnormal_mask, 1:] = disease_labels  # 更新类别1-7

    return final_pred

def collect_result(output: Dict):
    for k, v in output.items():
        result_cache[k].append(v.detach().cpu())
    return result_cache

def merge_result():
    collected = {}
    for k, v in result_cache.items():
        if len(v[0].shape) == 0:
            collected[k] = torch.stack(v)
        else:
            collected[k] = torch.cat(v)
    result_cache.clear()
    return collected

def metrics(data, num_classes=8) -> dict:
    # 忽略UndefinedMetricWarning警告
    warnings.filterwarnings("ignore", category=UserWarning, module="sklearn.metrics._classification")

    logits = data['predictions'].cpu().numpy()
    probabilities = torch.sigmoid(torch.tensor(logits)).numpy()

    threshold = 0.5
    predictions = (probabilities > threshold).astype(int)
    
    labels = data['labels'].cpu().numpy()
    #计算每个类别的准确率
    class_matrix = {}
    for i in range(num_classes):
        # 提取当前类别的预测标签和真实标签
        class_preds = predictions[:, i]
        class_labels = labels[:, i]

        accuracy = accuracy_score(class_labels, class_preds)
        precision = precision_score(class_labels, class_preds)
        recall = recall_score(class_labels, class_preds)
        # 统计样本数
        num = np.sum(class_labels)
        class_matrix[class_name[i]] = {
            "number" : int(num),
            "accuracy" : accuracy,
            "precision" : precision,
            "recall" : recall,
        }

    acc = accuracy_score(labels, predictions)
    precision = precision_score(labels, predictions, average='macro')
    recall = recall_score(labels, predictions, average='macro')
    auc = roc_auc_score(labels, probabilities, average='macro')

    result = dict(
        accuracy=acc,
        precision=precision,
        recall=recall,
        auc=auc,
        class_matrix=class_matrix
    )
    return result

def predict(model, dataset, mode):
    model.eval()
    model.to(device)
    print("Predicting...")
    with torch.no_grad():
        if mode == 'formal':
            for i_batch, data in enumerate(dataset):
                output = batch_formal(model=model,data=data)
                collect_result(output)  
            print('Double eyes test finished')
        elif mode == 'test':
            for i_batch, data in enumerate(dataset):
                output = batch(model=model,data=data)
                collect_result(output)      
            print('Single eye test finished')
        else:
            raise Exception   
    merged_output = merge_result()
    return merged_output

if __name__ == '__main__':
    # model_path = "logs/bs16_epoch50_lr0.0001/TrainerDR_20250208_thirdtry/model_049.pth"
    # model_path = "logs/bs16_epoch50_lr0.0001/TrainerDR_20250214_single1/model_049.pth"
    # model_path = "logs/bs16_epoch100_lr0.0001/TrainerDR_20250218_beta0.78/model_049.pth"
    cfg = TrainerConfig()

    model = DeepSurModel(backbone=cfg.model).to(device)
    if cfg.model != None:
        checkpoint = torch.load(cfg.model_path, map_location=device, weights_only=True)
        model.load_state_dict(checkpoint['model_state_dict'])
    trainer = TrainerDR()
    test_dataloader = trainer.test_loader
    formal_test_dataloader = trainer.formal_test_loader

    output_dir = os.path.join('predict_result', cfg.result_path)
    output_file_path1 = f'{cfg.result_path}_single.json'
    output_file_path2 = f'{cfg.result_path}_double.json'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    output_file1 = open(os.path.join(output_dir, output_file_path1), 'w', encoding='utf-8')
    output_file2 = open(os.path.join(output_dir, output_file_path2), 'w', encoding='utf-8')

    # 单眼测试
    output1 = predict(model, test_dataloader, mode='val')
    metrics1 = metrics(data=output1)
    print(json.dumps(metrics1, ensure_ascii=False), file=output_file1)
    output_file1.flush()

    # 双眼测试
    output2 = predict(model, formal_test_dataloader, mode='formal')
    metrics2 = metrics(data=output2)
    print(json.dumps(metrics2, ensure_ascii=False), file=output_file2)
    output_file2.flush()

    print(f"Result saved in {output_dir}")